from datetime import datetime, timedelta
from decimal import ROUND_HALF_UP, Decimal
import a_trade_calendar
from pandas import DataFrame, to_datetime

class CommonUtil:

    """
    股票常用帮助方法
    """
    @classmethod
    def get_period_startdate(cls, kline_period:str) :
        """
        根据当前周期的第一天
        :param kline_period: 周期
        :return : 当前周期的第一天
        """
        match(kline_period):
            case 'daily':
                return None 
            case 'weekly':
                today = datetime.now().date()
                monday = today - timedelta(days=today.weekday())
                return monday
            case 'monthly':
                firstday = datetime.now().replace(day=1).date()
                return firstday

    @classmethod
    def is_date_in_current_week(cls,target_date):
        """
        判断日期是否在当前周内
        :param target_date: 目标日期
        :return : 是否在当前周内
        """
        # 获取当前日期和目标日期的年、周数
        current_year, current_week, _ = datetime.now().isocalendar()
        target_year, target_week, _ = target_date.isocalendar()
        # 比较年份和周数
        return current_year == target_year and current_week == target_week
    
    @classmethod
    def is_date_in_current_month(cls,target_date):
        """
        判断日期是否在当前月内
        :param target_date: 目标日期
        :return : 是否在当前月内
        """
        now = datetime.now()
        return target_date.year == now.year and target_date.month == now.month
    
    @classmethod
    def get_latest_trade_date(cls):
        """
        获取最新交易日
        :return : 最新交易日
        """
        # 
        latest_trade_date_str = a_trade_calendar.get_latest_trade_date()
        format = "%Y-%m-%d"
        latest_trade_date = datetime.strptime(latest_trade_date_str, format).date()

        today_str = datetime.now().strftime(format)   
        hour = datetime.now().hour 
        is_today_trade = a_trade_calendar.is_trade_date(today_str)
        if is_today_trade :
            if hour > 15:
                return latest_trade_date
            else :
                return latest_trade_date - timedelta(days=1)
        else:
            return  latest_trade_date
    
class KlineUtil:
    """
    k线帮助方法
    """

    @classmethod
    def convert_to_custom_freq(cls, df_daily:DataFrame, freq='W') -> DataFrame:
        """
        将日线数据转换为指定频率的K线数据
        :param df_daily: 日线DataFrame
        :param freq: 目标频率，'W'为周，'M'为月
        :return: 转换后的DataFrame
        """
        agg_rules = {
            'trade_date': 'last',
            'symbol': 'last',
            'open': 'first',
            'close': 'last',
            'high': 'max',
            'low': 'min',
            'vol': 'sum',
            'amount': 'sum',
            'turn': 'sum'
        }
        df_daily.set_index('trade_date', drop=False, append=False, inplace=True)
        df_daily.index = to_datetime(df_daily.index)
        df_converted = df_daily.resample(freq).agg(agg_rules).dropna()
        
        # 重新计算涨跌额和涨跌幅、振幅
        df_converted['amplitude'] =  df_converted.apply(
            lambda x: ((x['high'] - x['low'])/x['low'] if x['low'] != 0 else 0) * 100, axis=1   
        )
        df_converted['amplitude'] = df_converted['amplitude'].map(
                lambda x: (Decimal(str(x)) * Decimal('1.00')) .quantize(Decimal('0.00'), rounding=ROUND_HALF_UP)
            )
        df_converted['chg'] = df_converted['close'] - df_converted['open']
        df_converted['chg'] = df_converted['chg'].map(
                lambda x: (Decimal(str(x)) * Decimal('1.00')).quantize(Decimal('0.00'), rounding=ROUND_HALF_UP)
            )
        df_converted['pct_chg'] = df_converted.apply(
            lambda x: (x['chg'] / x['open'] if x['open'] != 0 else 0) * 100 , axis=1   
        )
        df_converted['pct_chg'] = df_converted['pct_chg'].map(
                lambda x: (Decimal(str(x)) * Decimal('1.00')).quantize(Decimal('0.00'), rounding=ROUND_HALF_UP)
            ) 
        return df_converted.reset_index(drop=True)